"""
This is part of Human Robot Dynamics-Control, By Tourists.

Based on "OpenLoong Dynamics Control"
Transfer from mujoco to ISAAC SIM by IsaacLab

"""
import mujoco
import glfw
import numpy as np
from scipy.spatial.transform import Rotation

# from .. import openloong
from openloong.sim_interface import UIctr
from openloong.sim_interface import MJ_Interface
from openloong.common import PVT_Ctr
from openloong.common import DataLogger, DataBus, MotionState, LegState
from openloong.alg import Pin_KinDyn
from openloong.mathq.useful_math import eul2Rot, eigen2std
from openloong.alg import WBC_priority
from openloong.alg import MPC
from openloong.alg import nx, nu
from openloong.alg import GaitScheduler
from openloong.alg import FootPlacement
from openloong.alg import JoyStickInterpreter

dt = 0.001
dt_200Hz = 0.005
# glfw.init()
# options = mujoco.MjOption()
# options.maxgeom = 10000

# // MuJoCo load and compile model
error = "Could not load binary model"
model_path = "robot_models/scene.xml"
mj_model = mujoco.MjModel.from_xml_path(model_path)
mj_data = mujoco.MjData(mj_model)

# glfw.init()

def main():
    # // initialize classes
    uiController = UIctr(mj_model,mj_data)  ## UI control for Mujoco
    mj_interface = MJ_Interface(mj_model, mj_data) ## data interface for Mujoco
    kinDynSolver = Pin_KinDyn("robot_models/AzureLoong.urdf")   ## kinematics and dynamics solver
    RobotState = DataBus(kinDynSolver.model_nv)  ## data bus
    WBC_solv = WBC_priority(kinDynSolver.model_nv, 18, 22, 0.7, mj_model.opt.timestep)  ## WBC solver
    MPC_solv = MPC(dt_200Hz)   ## mpc controller
    gaitScheduler = GaitScheduler(0.25, mj_model.opt.timestep)   ##  gait scheduler
    pvtCtr = PVT_Ctr(mj_model.opt.timestep, "user_examples/openloong/common/joint_ctrl_config.json")  ## PVT joint control
    footPlacement = FootPlacement()  ## foot-placement planner
    jsInterp = JoyStickInterpreter(mj_model.opt.timestep)  ##desired baselink velocity generator
    logger = DataLogger("user_examples/record/datalog.log")  ## data logger

    # scene = mujoco.MjvScene(mj_model, maxgeom=20000)
    # mujoco.mjv_updateScene(
    #     mj_model, mj_data, mujoco.MjvOption(), None,
    #     mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL, scene)
    
    # // initialize UI: GLFW
    uiController.iniGLFW()
    uiController.enableTracking()  ##  enable viewpoint tracking of the body 1 of the robot
    uiController.createWindow("Demo", False)  ##  NOTE: if the saveVideo is set to true, the raw recorded file could be 2.5 GB for 15 seconds!

    # // initialize variables
    stand_legLength = 1.01  ## //0.97;// desired baselink height
    foot_height =0.07  ## distance between the foot ankel joint and the bottom
    xv_des = 1.2  ## desired velocity in x direction
    model_nv = kinDynSolver.model_nv

    RobotState.width_hips = 0.229
    footPlacement.kp_vx = 0.03
    footPlacement.kp_vy = 0.03
    footPlacement.kp_wz = 0.03
    footPlacement.stepHeight = 0.15
    footPlacement.legLength=stand_legLength

    # key_qpos_1d = mj_model.key_qpos.flatten()
    # mujoco.mju_copy(mj_data.qpos, mj_model.key_qpos, mj_model.nq * 1) ## set ini pos in Mujoco
    # Select the first row of key_qpos and reshape it
    # initial_pose = mj_model.key_qpos[0].reshape(-1, 1)
    
    # initial_pose = mj_model.key_qpos[0]
    # mj_data.qpos = initial_pose

    # Reshape mj_data.qpos to match the expected shape
    # mj_data.qpos = mj_data.qpos.reshape(-1, 1)

    # Now copy the initial pose to qpos
    # mujoco.mju_copy(mj_data.qpos, initial_pose, mj_model.nq)
    # mj_data.qpos = mj_data.qpos.flatten()
    
    # Optionally, reset the simulation to apply the new initial state
    # mujoco.mj_forward(mj_model, mj_data)

    motors_pos_des = np.zeros(model_nv - 6, dtype=np.float64)
    motors_pos_cur = np.zeros(model_nv - 6, dtype=np.float64)
    motors_vel_des = np.zeros(model_nv - 6, dtype=np.float64)
    motors_vel_cur = np.zeros(model_nv - 6, dtype=np.float64)
    motors_tau_des = np.zeros(model_nv - 6, dtype=np.float64)
    motors_tau_cur = np.zeros(model_nv - 6, dtype=np.float64)

    # // ini position and posture for foot-end and hand
    fe_l_pos_L_des = np.array([-0.018, 0.113, -stand_legLength])
    fe_r_pos_L_des = np.array([-0.018, -0.116, -stand_legLength])
    fe_l_eul_L_des = np.array([-0.000, -0.008, -0.000])
    fe_r_eul_L_des = np.array([0.000, -0.008, 0.000])
    fe_l_rot_des = eul2Rot(fe_l_eul_L_des[0], fe_l_eul_L_des[1], fe_l_eul_L_des[2])
    fe_r_rot_des = eul2Rot(fe_r_eul_L_des[0], fe_r_eul_L_des[1], fe_r_eul_L_des[2])

    hd_l_pos_L_des = np.array([-0.02, 0.32, -0.159])
    hd_r_pos_L_des = np.array([-0.02, -0.32, -0.159])
    hd_l_eul_L_des = np.array([-1.7581, 0.2129, 2.9581])
    hd_r_eul_L_des = np.array([1.7581, 0.2129, -2.9581])
    hd_l_rot_des = eul2Rot(hd_l_eul_L_des[0], hd_l_eul_L_des[1], hd_l_eul_L_des[2])
    hd_r_rot_des = eul2Rot(hd_r_eul_L_des[0], hd_r_eul_L_des[1], hd_r_eul_L_des[2])

    # Assuming kinDynSolver is an instance of a class with computeInK_Leg and computeInK_Hand methods
    resLeg = kinDynSolver.computeInK_Leg(fe_l_rot_des, fe_l_pos_L_des, fe_r_rot_des, fe_r_pos_L_des)
    resHand = kinDynSolver.computeInK_Hand(hd_l_rot_des, hd_l_pos_L_des, hd_r_rot_des, hd_r_pos_L_des)
    # Assuming mj_model is a MuJoCo model object
    qIniDes = np.zeros(mj_model.nq)
    qIniDes[7:] = resLeg.jointPosRes + resHand.jointPosRes
    # Assuming WBC_solv is an instance of a class with setQini method
    # and RobotState is an object with a q attribute
    WBC_solv.set_qini(qIniDes, RobotState.q)

    # // register variable name for data logger
    logger.add_item("simTime",1)
    logger.add_item("motor_pos_des",model_nv - 6)
    logger.add_item("motor_pos_cur",model_nv - 6)
    logger.add_item("motor_vel_des",model_nv - 6)
    logger.add_item("motor_vel_cur",model_nv - 6)
    logger.add_item("motor_tor_des",model_nv - 6)
    logger.add_item("rpyVal",3)
    logger.add_item("base_omega_W",3)
    logger.add_item("gpsVal",3)
    logger.add_item("base_vel",3)
    logger.add_item("dX_cal",12)
    logger.add_item("Ufe",12)
    logger.add_item("Xd",12)
    logger.add_item("X_cur",12)
    logger.add_item("X_cal",12)
    logger.finish_item_adding()

    """ -------------------------- main loop -------------------------------- """

    MPC_count = 0  ## count for controlling the mpc running period

    startSteppingTime = 3
    startWalkingTime = 5
    simEndTime = 30

    simstart = mj_data.time
    simTime = mj_data.time

    while not glfw.window_should_close(uiController.window):
        simstart = mj_data.time
        while (mj_data.time - simstart < 1.0 / 60.0 and uiController.runSim): ## press "1" to pause and resume, "2" to step the simulation
            mujoco.mj_step(mj_model, mj_data)
            simTime = mj_data.time
            #  Read the sensors:
            mj_interface.updateSensorValues()
            mj_interface.dataBusWrite(RobotState)

            ## update kinematics and dynamics info
            kinDynSolver.dataBusRead(RobotState)
            kinDynSolver.computeJ_dJ()
            kinDynSolver.computeDyn()
            kinDynSolver.dataBusWrite(RobotState)

            ## joint number: arm-l: 0-6, arm-r: 7-13, head: 14, waist: 15-17, leg-l: 18-23, leg-r: 24-29
            if (simTime > startWalkingTime):
                jsInterp.setWzDesLPara(0, 1)
                jsInterp.setVxDesLPara(xv_des, 2.0) ## jsInterp.setVxDesLPara(0.9,1);
                RobotState.motionState = MotionState.Walk ## start walking
            else:
                jsInterp.setIniPos(RobotState.q[0], RobotState.q[1], RobotState.base_rpy[2])
            jsInterp.step()
            RobotState.js_pos_des[2] = stand_legLength + foot_height  ## pos z is not assigned in jyInterp
            jsInterp.dataBusWrite(RobotState)  ## only pos x, pos y, theta z, vel x, vel y , omega z are rewrote.

            if (simTime >= startSteppingTime):
                ## gait scheduler
                gaitScheduler.dataBusRead(RobotState)
                gaitScheduler.step()
                gaitScheduler.dataBusWrite(RobotState)
                
                footPlacement.dataBusRead(RobotState)
                footPlacement.getSwingPos()
                footPlacement.dataBusWrite(RobotState)

            ## ------------- MPC ------------
            MPC_count = MPC_count + 1
            if MPC_count > (dt_200Hz / dt-1):
                MPC_solv.dataBusRead(RobotState)
                MPC_solv.cal()
                MPC_solv.dataBusWrite(RobotState)
                MPC_count = 0

            ## ------------- WBC ------------
            ## WBC Calculation
            WBC_solv.dataBusRead(RobotState)
            WBC_solv.computeDdq(kinDynSolver)
            WBC_solv.computeTau()
            WBC_solv.dataBusWrite(RobotState)
            ## get the final joint command
            if simTime <= startSteppingTime:
                RobotState.motors_pos_des = eigen2std(resLeg.jointPosRes + resHand.jointPosRes)
                RobotState.motors_vel_des = motors_vel_des
                RobotState.motors_tor_des = motors_tau_des
            else:
                MPC_solv.enable()
                nx = 12  # Example value, adjust as needed
                nu = 13  # Example value, adjust as needed

                # Create L_diag
                L_diag = np.array([
                    1.0, 1.0, 1.0,    # eul
                    1.0, 200.0, 1.0,  # pCoM
                    1e-7, 1e-7, 1e-7, # w
                    100.0, 10.0, 1.0  # vCoM
                ]).reshape(1, nx)

                # Create K_diag
                K_diag = np.array([
                    1.0, 1.0, 1.0,  # fl
                    1.0, 1.0, 1.0,
                    1.0, 1.0, 1.0,  # fr
                    1.0, 1.0, 1.0, 1.0
                ]).reshape(1, nu)

                # Assuming MPC_solv is an instance of a class with set_weight method
                MPC_solv.set_weight(1e-6, L_diag, K_diag)

                # Integrate and set desired positions, velocities, and torques
                pos_des = kinDynSolver.integrateDIY(RobotState.q, RobotState.wbc_delta_q_final)
                RobotState.motors_pos_des = eigen2std(pos_des[7:7+model_nv-6])
                RobotState.motors_vel_des = eigen2std(RobotState.wbc_dq_final)
                RobotState.motors_tor_des = eigen2std(RobotState.wbc_tauJointRes)

                # Joint PVT controller
                pvtCtr.dataBusRead(RobotState)
                if simTime <= 3:
                    pvtCtr.calMotorsPVT(100.0 / 1000.0 / 180.0 * 3.1415)
                else:
                    pvtCtr.setJointPD(100, 10, "J_ankle_l_pitch")
                    pvtCtr.setJointPD(100, 10, "J_ankle_l_roll")
                    pvtCtr.setJointPD(100, 10, "J_ankle_r_pitch")
                    pvtCtr.setJointPD(100, 10, "J_ankle_r_roll")
                    pvtCtr.setJointPD(1000, 100, "J_knee_l_pitch")
                    pvtCtr.setJointPD(1000, 100, "J_knee_r_pitch")
                    pvtCtr.calMotorsPVT()
                pvtCtr.dataBusWrite(RobotState)

                # Give the joint torque command to the interface
                mj_interface.setMotorsTorque(RobotState.motors_tor_out)

            ## print info to the console
            # printf("f_L=[%.3f, %.3f, %.3f]\n", RobotState.fL[0], RobotState.fL[1], RobotState.fL[2]);
            # printf("f_R=[%.3f, %.3f, %.3f]\n", RobotState.fR[0], RobotState.fR[1], RobotState.fR[2]);

            # printf("rpyVal=[%.5f, %.5f, %.5f]\n", RobotState.rpy[0], RobotState.rpy[1], RobotState.rpy[2]);
            # printf("basePos=[%.5f, %.5f, %.5f]\n", RobotState.basePos[0], RobotState.basePos[1], RobotState.basePos[2]);

            ## data save
            logger.start_new_line()
            logger.record_item_data("simTime", simTime)
            logger.record_item_data("motor_pos_des", RobotState.motors_pos_des)
            logger.record_item_data("motor_pos_cur", RobotState.motors_pos_cur)
            logger.record_item_data("motor_vel_des", RobotState.motors_vel_des)
            logger.record_item_data("motor_vel_cur", RobotState.motors_vel_cur)
            logger.record_item_data("motor_tor_des", RobotState.motors_tor_des)
            logger.record_item_data("rpyVal", RobotState.rpy)
            logger.record_item_data("base_omega_W", RobotState.base_omega_W)
            logger.record_item_data("gpsVal", RobotState.basePos)
            logger.record_item_data("base_vel", RobotState.dq[:3])
            logger.record_item_data("dX_cal",RobotState.dX_cal)
            logger.record_item_data("Ufe",RobotState.Fr_ff)
            logger.record_item_data("Xd",RobotState.Xd)
            logger.record_item_data("X_cur",RobotState.X_cur)
            logger.record_item_data("X_cal",RobotState.X_cal)
            logger.finish_line()

        if mj_data.time >= simEndTime:
            break
        
        
        # scene = mujoco.MjvScene(mj_model, maxgeom=20000)
        # mujoco.mjv_updateScene(
        #     mj_model, mj_data, mujoco.MjvOption(), None,
        #     mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL, scene)
        uiController.updateScene()
    ## free visualization storage
    uiController.close()

    return 0


if __name__ == "__main__":
    main()